namespace NBody
{
    using System;
    using System.Collections;
    using System.IO;
    using System.Runtime.InteropServices;
    
    public class NBodySystem : IEnumerable, ICloneable
    {
        private int numparts;
        private double snaptime;
        private Particle[] particle;
        
        private int size;
        private bool is_sorted_by_radius;
        
        public NBodySystem(int size, double snaptime)
        {
            this.size = size;
            this.snaptime = snaptime;
            this.particle = new Particle[size];
            
            this.numparts = 0;
            
            is_sorted_by_radius = false;
        }
        
        public int NumParts { get { return numparts; } }
        
        public double Time { get { return snaptime; } }
        public void SetTime(double t)
        {
            snaptime = t;
        }
        
        public int Size { get { return size; } }
        
        public bool IsSortedByRadius { get { return is_sorted_by_radius; } }
        
        public Particle Particle(int i)
        {
            return particle[i];
        }
        
        public void AddParticle(Particle p)
        {
            if (numparts < size)
            {
                particle[numparts++] = p;
            }
            else
            {
                throw new IndexOutOfRangeException
                          ("Cannot add another particle.");
            }
            
            is_sorted_by_radius = false;
        }
        
        // removes ith particle.  Does not preserve particle order!
        public void RemoveParticle(int i)
        {
            particle[i] = particle[numparts - 1];
            numparts = numparts - 1;
            is_sorted_by_radius = false;
        }
        
        public Particle this [int i]
        {
            get { return particle[i]; }
            set
            {
                particle[i] = value;
                is_sorted_by_radius = false;
            }
        }
        
        public IEnumerator GetEnumerator()
        {
            return new SystemEnumerator(this);
        }
        
        public double Mass
        {
            get
            {
                double m = 0;
                for (int i = 0; i < numparts; i++)
                    m += particle[i].Mass;
                
                return m;
            }
        }
        
        public double Rmax
        {
            get
            {
                double rmax = particle[0].Radius;
                for (int i = 1; i < numparts; i++)
                {
                    double test = particle[i].Radius;
                    if (test > rmax) rmax = test;
                }
                
                return rmax;
            }
        }
        
        public double Rmin
        {
            get
            {
                double rmin = particle[0].Radius;
                for (int i = 1; i < numparts; i++)
                {
                    double test = particle[i].Radius;
                    if (test < rmin) rmin = test;
                }

                return rmin;
            }
        }
        
        // move origin to coordinates given
        // doesn't change velocities!!!
        public void TranslateTo(double new_x, double new_y, double new_z)
        {
            Console.Error.WriteLine("Translating system to new origin ({0:f6}, {1:f6}, {2:f6})", new_x, new_y, new_z);
            
            for (int i = 0; i < numparts; i++)
            {
                double x = particle[i].X - new_x;
                double y = particle[i].Y - new_y;
                double z = particle[i].Z - new_z;
                particle[i].SetPosition((double)x, (double)y, (double)z);
            }
        }
        
        public void RotateAboutX(double degrees)
        {
            Console.Error.WriteLine("Rotating system by {0} degrees about X axis", degrees);
            double theta = degrees / 180.0 * Math.PI;
            
            double a, b, va, vb;
            
            for(int i = 0; i < numparts; i++)
            {
                a = particle[i].Y * Math.Cos(theta) + particle[i].Z * Math.Sin(theta);
                b = -particle[i].Y * Math.Sin(theta) + particle[i].Z * Math.Cos(theta);
                va = particle[i].Vy * Math.Cos(theta) + particle[i].Vz * Math.Sin(theta);
                vb = -particle[i].Vy * Math.Sin(theta) + particle[i].Vz * Math.Cos(theta);
        
                particle[i].SetPosition(particle[i].X, (double)a, (double)b);
                particle[i].SetVelocity(particle[i].Vx, (double)va, (double)vb);
            }
        }
        
        public void RotateAboutY(double degrees)
        {
            Console.Error.WriteLine("Rotating system by {0} degrees about Y axis", degrees);
            double theta = degrees / 180.0 * Math.PI;
            
            double a, b, va, vb;
            
            for(int i = 0; i < numparts; i++)
            {
                a = particle[i].X * Math.Cos(theta) + particle[i].Z * Math.Sin(theta);
                b = -particle[i].X * Math.Sin(theta) + particle[i].Z * Math.Cos(theta);
                va = particle[i].Vx * Math.Cos(theta) + particle[i].Vz * Math.Sin(theta);
                vb = -particle[i].Vx * Math.Sin(theta) + particle[i].Vz * Math.Cos(theta);
        
                particle[i].SetPosition((double)a, particle[i].Y, (double)b);
                particle[i].SetVelocity((double)va, particle[i].Vy, (double)vb);
            }
        }
        
        public void RotateAboutZ(double degrees)
        {
            Console.Error.WriteLine("Rotating system by {0} degrees about Z axis", degrees);
            double theta = degrees / 180.0 * Math.PI;
            
            double a, b, va, vb;
            
            for(int i = 0; i < numparts; i++)
            {
                a = particle[i].X * Math.Cos(theta) + particle[i].Y * Math.Sin(theta);
                b = -particle[i].X * Math.Sin(theta) + particle[i].Y * Math.Cos(theta);
                va = particle[i].Vx * Math.Cos(theta) + particle[i].Vy * Math.Sin(theta);
                vb = -particle[i].Vx * Math.Sin(theta) + particle[i].Vy * Math.Cos(theta);
        
                particle[i].SetPosition((double)a, (double)b, particle[i].Z);
                particle[i].SetVelocity((double)va, (double)vb, particle[i].Vz);
            }
        }
        
        public void RotateByMatrix(double[,] m)
        {
            Console.Error.WriteLine("Rotating system by matrix");
            
            double x, y, z, vx, vy, vz;
            for (int i = 0; i < numparts; i++)
            {
                x = m[0, 0] * particle[i].X + m[0, 1] * particle[i].Y + m[0, 2] * particle[i].Z;
                y = m[1, 0] * particle[i].X + m[1, 1] * particle[i].Y + m[1, 2] * particle[i].Z;
                z = m[2, 0] * particle[i].X + m[2, 1] * particle[i].Y + m[2, 2] * particle[i].Z;
                
                vx = m[0, 0] * particle[i].Vx + m[0, 1] * particle[i].Vy + m[0, 2] * particle[i].Vz;
                vy = m[1, 0] * particle[i].Vx + m[1, 1] * particle[i].Vy + m[1, 2] * particle[i].Vz;
                vz = m[2, 0] * particle[i].Vx + m[2, 1] * particle[i].Vy + m[2, 2] * particle[i].Vz;
                
                particle[i].SetPosition(x, y, z);
                particle[i].SetVelocity(vx, vy, vz);
            }
        }
        
        public void AlignWithVector(double x, double y, double z)
        {
            Console.Error.WriteLine("Aligning system with vector ({0}, {1}, {2})", x, y, z);
            
            double r = Math.Sqrt(x*x + y*y + z*z);
            double theta_z = 180.0 / Math.PI * Math.Atan2(y, x);
            double theta_x = 180.0 / Math.PI * Math.Acos(z / r);
        
            this.RotateAboutZ(theta_z);
            this.RotateAboutY(-theta_x);
        }
        
        public void FindCentreOfMass(out double x, out double y, out double z)
        {
            x = y = z = 0.0;
            double mtot = 0.0;
            for (int i = 0; i < numparts; i++)
            {
                x += particle[i].Mass * particle[i].X;
                y += particle[i].Mass * particle[i].Y;
                z += particle[i].Mass * particle[i].Z;
                mtot += particle[i].Mass;
            }
            x /= mtot;
            y /= mtot;
            z /= mtot;
            
            Console.Error.WriteLine("Global c/m coordinates: {0:f6}, {1:f6}, {2:f6}", x, y, z);
        }
        
        public void FindCentreOfMass(double radius, double[] centre, out double[] cm, double tolerance, out double mtot)
        {
            cm = new double[] {0.0, 0.0, 0.0};
            double[] cm_last = new double[] {centre[0], centre[1], centre[2]};
            int count = 0;
            double error;
            do
            {
                double x, y, z, r;
                mtot = 0.0;
                cm[0] = cm[1] = cm[2] = 0.0;
                for (int i = 0; i < numparts; i++)
                {
                    x = particle[i].X - cm_last[0];
                    y = particle[i].Y - cm_last[1];
                    z = particle[i].Z - cm_last[2];
                    r = Math.Sqrt(x*x + y*y + z*z);
                    if (r < radius)
                    {
                        mtot += particle[i].Mass;
                        cm[0] += particle[i].Mass * x;
                        cm[1] += particle[i].Mass * y;
                        cm[2] += particle[i].Mass * z;
                    }
                }
                cm[0] /= mtot;
                cm[1] /= mtot;
                cm[2] /= mtot;                
                               
                // move to original coordinate system
                cm[0] += cm_last[0];
                cm[1] += cm_last[1];
                cm[2] += cm_last[2];
                
                double error_x = cm[0] - cm_last[0];
                double error_y = cm[1] - cm_last[1];
                double error_z = cm[2] - cm_last[2];
                error = Math.Sqrt(error_x*error_x + error_y*error_y + error_z*error_z);
                
                cm_last[0] = cm[0];
                cm_last[1] = cm[1];
                cm_last[2] = cm[2];
                
                if (count++ > 50)
                    throw new NBodyException("FindCentreOfMass won't converge after 50 iterations");
            }
            while (error > tolerance);
            //Console.Error.WriteLine("c/m coordinates: {0:f6}, {1:f6}, {2:f6}", cm[0], cm[1], cm[2]);
        }    
        
        public void FindCentreOfMass(out double x, out double y, out double z, double tolerance)
        {
            // use global c/m as initial guess
            double[] cm = new double[3];
            double[] new_cm = new double[3];
            FindCentreOfMass(out cm[0], out cm[1], out cm[2]);
            
            double mtot = Mass;
            double rmax = Rmax;
            bool done = false;
            while (!done)
            {
                // shrink radius by ... 25%??
                rmax *= 0.75;
                Console.Error.WriteLine("  Finding c/m for radius of {0}", rmax);
                double m;
                FindCentreOfMass(rmax, cm, out new_cm, 0.01 * rmax, out m);
                if (m < 0.01 * mtot)
                {
                    Console.Error.WriteLine("    Enclosed mass is now less than 1% of total mass; stopping");
                    done = true;
                }
                
                double error_x = new_cm[0] - cm[0];
                double error_y = new_cm[1] - cm[1];
                double error_z = new_cm[2] - cm[2];
                double error = Math.Sqrt(error_x*error_x + error_y*error_y + error_z*error_z);
                //Console.Error.WriteLine("    Error is {0}", error);
                if (error < tolerance && m < 0.5 * mtot)
                    done = true;
                    
                cm[0] = new_cm[0];
                cm[1] = new_cm[1];
                cm[2] = new_cm[2];
                //Console.Error.WriteLine("r = {0}: {1}, {2}, {3}", rmax, cm[0], cm[1], cm[2]);
            }
            
            x = new_cm[0];
            y = new_cm[1];
            z = new_cm[2];
            Console.Error.WriteLine("Final c/m coordinates: {0:f6}, {1:f6}, {2:f6}", x, y, z);
        }
        
        /*public void FindCentreOfMass(out double x, out double y, out double z, double tolerance)
        {
            double[] rcm = new double[3];
            FindCentreOfMass(out rcm[0], out rcm[1], out rcm[2]);
            
            Console.Error.WriteLine("  Iterating for a better estimate ...");
            double[] rcm_last = {rcm[0], rcm[1], rcm[2]};
            
            // find maximum radius of system centered on c/m
            double xx = particle[0].X - rcm[0];
            double yy = particle[0].Y - rcm[1];
            double zz = particle[0].Z - rcm[2];
            double rmax = Math.Sqrt(xx * xx + yy * yy + zz * zz);
            for (int i = 1; i < numparts; i++)
            {
                xx = particle[i].X - rcm[0];
                yy = particle[i].Y - rcm[1];
                zz = particle[i].Z - rcm[2];
                double test = Math.Sqrt(xx * xx + yy * yy + zz * zz);
                if (test > rmax) rmax = test;
            }
            
            double ri = rmax;
            bool done = false;
            int c = 0;
            while (!done)
            {
                ri = 0.8 * ri;
                double change = Double.PositiveInfinity;
                c++;
                while (change > tolerance)
                {
                    // find c/m of all particles within ri
                    int n_inside = 0;
                    double total_mass = 0.0;
                    foreach (Particle p in particle)
                    {
                        xx = p.X - rcm_last[0];
                        yy = p.Y - rcm_last[1];
                        zz = p.Z - rcm_last[2];
                        if (Math.Sqrt(xx * xx + yy * yy + zz * zz) <= ri)
                        {
                            for (int i = 0; i < 3; i++)
                                rcm[i] += p.Mass * p.Position(i);
                            total_mass += p.Mass;
                            n_inside++;
                        }
                    }
                    Console.Error.WriteLine("{0} {1} {2}", rcm[0], rcm[1], rcm[2]);
                    for (int i = 0; i < 3; i++)
                        rcm[i] /= total_mass;
                    
                    // keep making radius smaller until there's
                    // less than 5% of the particles inside
                    if ((double)n_inside < 0.05 * (double)numparts)
                    {
                        done = true;
                        break;
                    }
                    
                    // check for convergence
                    change = Math.Abs(rcm[0] - rcm_last[0]);
                    for (int i = 1; i < 3; i++)
                    {
                        if (Math.Abs(rcm[i] - rcm_last[i]) > change)
                            change = Math.Abs(rcm[i] - rcm_last[i]);
                    }
                    for (int i = 0; i < 3; i++)
                        rcm_last[i] = rcm[i];
                    
                }
                //Console.Error.WriteLine("{0:f6}, {1:f6}, {2:f6}", rcm[0], rcm[1], rcm[2]);
            }
            //Console.Error.WriteLine("count = {0}", c);
            x = rcm[0];
            y = rcm[1];
            z = rcm[2];            
            Console.Error.WriteLine("  Final c/m coords: {0:f6}, {1:f6}, {2:f6}", x, y, z);
        }*/
        
        public void AdjustForCentreOfMass()
        {
            double x, y, z;
            FindCentreOfMass(out x, out y, out z);
            TranslateTo((double)x, (double)y, (double)z);
        }
        
        public void AdjustForCentreOfMass(double tolerance)
        {
            double x, y, z;
            FindCentreOfMass(out x, out y, out z, tolerance);
            TranslateTo(x, y, z);
        }
        
        // "Remove" all particles beyond a certain radius
        public void ExtractSphere(double radius)
        {
            Console.Error.WriteLine("Removing all particles beyond a radius of {0}", radius);
            int i;
            for (i = 0; i < numparts; i++)
            {
                if (particle[i].Radius > radius)
                {
                    RemoveParticle(i);
                    i--;
                }
            }
        }
        
        // "Remove" all particles beyond a certain cylindrical radius
        public void ExtractCylinder(double radius, double height)
        {
            Console.Error.WriteLine("Removing all particles beyond a cylindrical radius of {0} and height {1}", radius, height);
            int i;
            for (i = 0; i < numparts; i++)
            {
                if (particle[i].CylindricalRadius > radius || Math.Abs(particle[i].Z) > height)
                {
                    RemoveParticle(i);
                    i--;
                }
            }
        }
        
        // very slow
        public void Sort(IComparer comp)
        {
            Console.Error.Write("Sorting system by custom IComparer ... ");
            Array.Sort(particle, comp);
            Console.Error.WriteLine("done");
            is_sorted_by_radius = false;
        }
        
        // by radius, built in for speed
        public void Sort()
        {
            Console.Error.Write("Sorting system by radius ... ");
            if (is_sorted_by_radius)
                Console.Error.WriteLine("already done");
            else
            {
                SortByRadius(0, numparts);
                Console.Error.WriteLine("done");
                is_sorted_by_radius = true;
            }
        }
        
        private void SortByRadius(int first, int n)
        {
            int pivot_index;
            int n1, n2;
    
            if (n > 1)
            {
                // partition inline
                Particle swap = particle[first + (n-1)/2]; // "middle" element
                particle[first + (n-1)/2] = particle[first];
                particle[first] = swap;
                Particle pivot = particle[first];
                double pivot_rad = pivot.RadiusSquared;
                int too_big_index = first + 1;
                int too_small_index = first + n - 1;
                
                while (too_big_index <= too_small_index)
                {
                    while (too_big_index < first + n && particle[too_big_index].RadiusSquared <= pivot_rad)
                        too_big_index++; 
                    while (particle[too_small_index].RadiusSquared > pivot_rad)
                        too_small_index--;
                    if (too_big_index < too_small_index)
                    {
                        swap = particle[too_big_index];
                        particle[too_big_index] = particle[too_small_index];
                        particle[too_small_index] = swap;
                    }
                }

                particle[first] = particle[too_small_index];
                particle[too_small_index] = pivot;
            
                pivot_index = too_small_index;
                // end inline partition
    
                n1 = pivot_index - first;
                n2 = n - n1 - 1;
        
                SortByRadius(first, n1);
                SortByRadius(pivot_index + 1, n2);
            }
        }
        
        public override string ToString()
        {
            string outs = String.Format("{0}\n{1:f6}\n", numparts, snaptime);
            for (int i = 0; i < numparts - 1; i++)
            {
                outs += particle[i] + "\n";
            }
            outs += particle[numparts - 1];
            return outs;
        }
        
        public static NBodySystem Read(string filename)
        {
            int n;
            double t;
            byte[] bytes;
            FileInfo finfo = new FileInfo(filename);
            
            using (BinaryReader br = new BinaryReader(finfo.OpenRead()))
            {
                n = br.ReadInt32();
                t = br.ReadDouble();
                
                // do some basic checking to make sure our file is alright ...
                // make sure total bytes of file match what the header tells us
                int total = 4 + 8 + n * 7 * 4;
                if (total != finfo.Length)
                {
                    string message = String.Format("Cannot read N-body file: mismatch between header ({0}) and size of file ({1})", total, finfo.Length);
                    throw new NBodyException(message);
                }
                
                // just in case, check to make sure that the particle number is sane
                if (n > 5000000)
                    throw new NBodyException("Are you sure you want to open a file with " + n.ToString() + " particles?");
                
                bytes = br.ReadBytes(n * 7 * 4);
            }
            
            Console.Error.WriteLine("Reading N-body system:\n  Number of Particles = {0}\n  Time = {1:f6}", n, t);
            
            NBodySystem s = new NBodySystem(n, t);
                        
            for (int i = 0; i < n; i++)
            {
                double m = BitConverter.ToSingle(bytes, 28 * i);
                double x = BitConverter.ToSingle(bytes, 28 * i + 4);
                double y = BitConverter.ToSingle(bytes, 28 * i + 8);
                double z = BitConverter.ToSingle(bytes, 28 * i + 12);
                double vx = BitConverter.ToSingle(bytes, 28 * i + 16);
                double vy = BitConverter.ToSingle(bytes, 28 * i + 20);
                double vz = BitConverter.ToSingle(bytes, 28 * i + 24);
                
                Particle p = new Particle(m, x, y, z, vx, vy, vz);
                s.AddParticle(p);
            }
            
            Console.Error.WriteLine("... done");
            
            return s;                        
        }
        
        public void Write(string filename)
        {   
            Console.Error.WriteLine("Writing N-body system:\n  Number of Particles = {0}\n  Time = {1:f6}", numparts, snaptime);
            
            byte[] buffer = new byte[28 * numparts];
            for (int i = 0; i < numparts; i++)
            {
                byte[] m = BitConverter.GetBytes((float)particle[i].Mass);
                byte[] x = BitConverter.GetBytes((float)particle[i].X);
                byte[] y = BitConverter.GetBytes((float)particle[i].Y);
                byte[] z = BitConverter.GetBytes((float)particle[i].Z);
                byte[] vx = BitConverter.GetBytes((float)particle[i].Vx);
                byte[] vy = BitConverter.GetBytes((float)particle[i].Vy);
                byte[] vz = BitConverter.GetBytes((float)particle[i].Vz);
                
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j] = m[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 4] = x[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 8] = y[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 12] = z[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 16] = vx[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 20] = vy[j];
                for (int j = 0; j < 4; j++)
                    buffer[28 * i + j + 24] = vz[j];
            }            
            
            using (BinaryWriter bw = new BinaryWriter(File.Open(filename, FileMode.Create)))
            {
                bw.Write(numparts);
                bw.Write(snaptime);
                bw.Write(buffer);                
            }
            
            Console.Error.WriteLine("... done");            
        }
        
        public object Clone()
        {
            NBodySystem clone = new NBodySystem(numparts, snaptime);
            for (int i = 0; i < numparts; i++)
                clone.AddParticle(new Particle(particle[i].Mass, particle[i].X, particle[i].Y, particle[i].Z, particle[i].Vx, particle[i].Vy, particle[i].Vz));
            
            return clone;
        }
        
        public class SystemEnumerator : IEnumerator
        {
            private int index;
            private NBodySystem system;
            
            public SystemEnumerator(NBodySystem s)
            {
                index = -1;
                system = s;
            }
            
            public void Reset()
            {
                index = -1;
            }
            
            public object Current
            {
                get
                {
                    if (index > -1)
                        return system[index];
                    else
                        return -1;
                }
            }
            
            public bool MoveNext()
            {
                index++;
                if (index < system.NumParts)
                    return true;
                else
                    return false;
            }
        }
        
        private static void check_particle_number(int n)
        {
            if (n > 3000000)
                throw new NBodyException("Are you sure you want to open a file with " + n.ToString() + " particles?");
        }
    }
}
